import asyncio
from asyncio import PriorityQueue, TimeoutError
from typing import Optional

from bald_spider.utils.pqueue import SpiderPriorityQueue


class Scheduler:
    """爬虫调度，用得是优先级队列"""

    def __init__(self):
        self.request_queue: Optional[SpiderPriorityQueue] = None

    def open(self):
        self.request_queue = SpiderPriorityQueue()

    async def next_request(self):
        request = await self.request_queue.get()
        return request

    async def enqueue_request(self, request):
        await self.request_queue.put(request)
